# encoding=utf-8
"""
需要这两个环境变量, 在pycharm命令行启动前设置
source /usr/local/Ascend/ascend-toolkit/set_env.sh
export PYTHONPATH=/usr/local/Ascend/thirdpart/aarch64/acllite:$PYTHONPATH
"""
import time
import cv2
import numpy
from AclTools import AclLiteResource
from acllite_model import AclLiteModel


class DensenetGrade:
    def __init__(self, pth_path: str):
        self.acl_resource = AclLiteResource()
        self.acl_resource.init()
        self.model = AclLiteModel(pth_path)

    def softmax(self, x):
        f_x = numpy.exp(x) / numpy.sum(numpy.exp(x))
        return f_x

    def infer(self, image_path: str):
        trans_dim = self._preprocess(image_path)
        result = self.model.execute([trans_dim, ])
        y = self.softmax(result[0][0]).tolist()
        return self._postprocess(y)

    def _postprocess(self, pred_arr):
        return {
            'index': pred_arr.index(max(pred_arr)),
            'pie': pred_arr
        }

    def _preprocess(self, data_in: str):
        img_arr = self.circle_crop(data_in) / 255.0
        x = numpy.expand_dims(img_arr[:, :, ::-1].astype(numpy.float32).transpose((2, 0, 1)), axis=0)
        return x

    def circle_crop(self, image_src: str):
        crop_mask = self.crop_image_from_mask(image_src)
        return self.crop_image_with_gaussian(crop_mask)

    def crop_image_from_mask(self, image_src: str):
        # load
        image = cv2.imread(image_src)

        # binary
        gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        _, binary_image = cv2.threshold(gray_image, 7, 255, cv2.THRESH_BINARY)

        # cal roi
        x, y, w, h = cv2.boundingRect(binary_image)
        center = (w // 2), (h // 2)
        radius = min(center)
        y = y + center[1] - radius
        x = x + center[0] - radius
        copy_image = image[y: y + 2 * radius, x: x + 2 * radius]

        # gen mask
        mask = numpy.zeros_like(copy_image)
        cv2.circle(mask, (radius, radius), radius, (1, 1, 1), -1)

        # exposure
        return copy_image * mask

    def crop_image_with_gaussian(self, data_in: numpy.ndarray):
        ori_image = cv2.resize(data_in, (224, 224)).astype(numpy.float32)
        out = cv2.GaussianBlur(ori_image, (11, 11), 10.0)
        exposure = cv2.addWeighted(ori_image, 4, out, -4, 128)
        exposure = numpy.clip(exposure, 0, 255).astype(numpy.uint8)
        exposure = cv2.cvtColor(exposure, cv2.COLOR_BGR2RGB)
        return exposure


if __name__ == '__main__':
    # init
    # init cost 1.6198961599984614
    t1 = time.perf_counter()
    grade = DensenetGrade("dense_om_quality.om")
    t2 = time.perf_counter()
    print("init cost", t2 - t1)

    # time cost
    # infer cost 0.13198743400062085
    # {'index': 2, 'pie': [0.147410586476326, 0.060305479913949966, 0.6969325542449951, 0.09030850231647491, 0.005042881704866886]}
    t1 = time.perf_counter()
    print(grade.infer("1.jpg"))
    t2 = time.perf_counter()
    print("infer cost", t2 - t1)

    # query resource
    import os
    os.system("npu-smi info")

    # release
    del grade
    time.sleep(3)
    # 不清楚是否释放成功
    os.system("npu-smi info")
